package main

import (
	"encoding/base64"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os/exec"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/pion/webrtc/v3"
	"github.com/pion/webrtc/v3/pkg/media"
)

type Signal struct {
	SDP       string `json:"sdp"`
	Type      string `json:"type"`
	Candidate string `json:"candidate,omitempty"`
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.Println("服务器启动...")
}

func main() {
	http.HandleFunc("/ws", handleWebSocket)
	log.Println("在 :8080 启动 HTTP 服务器")
	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatalf("HTTP 服务器启动失败: %v", err)
	}
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	log.Println("收到来自", r.RemoteAddr, "的 WebSocket 连接请求")
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket 升级失败: %v", err)
		return
	}
	defer func() {
		log.Println("关闭 WebSocket 连接")
		conn.Close()
	}()

	// WebRTC 配置
	log.Println("使用 STUN 初始化 WebRTC 配置")
	config := webrtc.Configuration{
		ICEServers: []webrtc.ICEServer{
			{URLs: []string{"stun:stun.l.google.com:19302"}},
		},
	}

	// 创建媒体引擎并注册 H.264 和 Opus 编解码器
	log.Println("创建媒体引擎并注册编解码器")
	mediaEngine := &webrtc.MediaEngine{}
	if err := mediaEngine.RegisterCodec(webrtc.RTPCodecParameters{
		RTPCodecCapability: webrtc.RTPCodecCapability{
			MimeType:    webrtc.MimeTypeH264,
			ClockRate:   90000,
			SDPFmtpLine: "level-asymmetry-allowed=1;packetization-mode=1;profile-level-id=42e01f",
		},
		PayloadType: 96,
	}, webrtc.RTPCodecTypeVideo); err != nil {
		log.Fatalf("注册 H.264 编解码器失败: %v", err)
	}
	if err := mediaEngine.RegisterCodec(webrtc.RTPCodecParameters{
		RTPCodecCapability: webrtc.RTPCodecCapability{
			MimeType:    webrtc.MimeTypeOpus,
			ClockRate:   48000,
			Channels:    2,
			SDPFmtpLine: "minptime=10;useinbandfec=1",
		},
		PayloadType: 111,
	}, webrtc.RTPCodecTypeAudio); err != nil {
		log.Fatalf("注册 Opus 编解码器失败: %v", err)
	}

	// 创建 API 和 PeerConnection
	log.Println("创建 WebRTC API 和 PeerConnection")
	api := webrtc.NewAPI(webrtc.WithMediaEngine(mediaEngine))
	peerConnection, err := api.NewPeerConnection(config)
	if err != nil {
		log.Fatalf("创建 PeerConnection 失败: %v", err)
	}
	defer func() {
		log.Println("关闭 PeerConnection")
		peerConnection.Close()
	}()

	// 创建视频轨道
	log.Println("创建视频轨道")
	videoTrack, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{
		MimeType:  webrtc.MimeTypeH264,
		ClockRate: 90000,
	}, "video", "pion-video")
	if err != nil {
		log.Fatalf("创建视频轨道失败: %v", err)
	}

	// 创建音频轨道
	log.Println("创建音频轨道")
	audioTrack, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{
		MimeType:  webrtc.MimeTypeOpus,
		ClockRate: 48000,
		Channels:  2,
	}, "audio", "pion-audio")
	if err != nil {
		log.Fatalf("创建音频轨道失败: %v", err)
	}

	// 将轨道添加到 PeerConnection
	log.Println("将视频轨道添加到 PeerConnection")
	if _, err := peerConnection.AddTrack(videoTrack); err != nil {
		log.Fatalf("添加视频轨道失败: %v", err)
	}
	log.Println("将音频轨道添加到 PeerConnection")
	if _, err := peerConnection.AddTrack(audioTrack); err != nil {
		log.Fatalf("添加音频轨道失败: %v", err)
	}

	// 处理 RTCP 反馈
	go func() {
		log.Println("启动 RTCP 反馈循环")
		rtcpBuf := make([]byte, 1500)
		for {
			n, _, err := peerConnection.GetTransceivers()[0].Receiver().Read(rtcpBuf)
			if err != nil {
				log.Printf("RTCP 读取错误: %v", err)
				return
			}
			log.Printf("收到 %d 字节 RTCP 数据", n)
		}
	}()

	// 缓存 ICE 候选者
	pendingCandidates := make([]webrtc.ICECandidateInit, 0)
	var remoteDescSet bool

	// 处理 ICE 候选者
	peerConnection.OnICECandidate(func(c *webrtc.ICECandidate) {
		if c != nil {
			log.Printf("新的 ICE 候选者: %s", c.ToJSON().Candidate)
			candidate, _ := json.Marshal(c.ToJSON())
			if err := conn.WriteJSON(Signal{
				Candidate: string(candidate),
				Type:      "candidate",
			}); err != nil {
				log.Printf("发送 ICE 候选者失败: %v", err)
			}
		}
	})

	// 处理 ICE 连接状态变化
	peerConnection.OnICEConnectionStateChange(func(state webrtc.ICEConnectionState) {
		log.Printf("ICE 连接状态改变: %s", state.String())
		if state == webrtc.ICEConnectionStateFailed {
			log.Println("ICE 连接失败，关闭 PeerConnection")
			peerConnection.Close()
			conn.WriteJSON(Signal{Type: "error", SDP: "ICE 连接失败"})
		}
	})

	// 处理 WebSocket 消息
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		log.Println("启动 WebSocket 消息循环")
		for {
			_, msg, err := conn.ReadMessage()
			if err != nil {
				log.Printf("WebSocket 读取错误: %v", err)
				return
			}
			log.Printf("收到 WebSocket 消息: %s", string(msg))

			var signal Signal
			if err := json.Unmarshal(msg, &signal); err != nil {
				log.Printf("JSON 解析错误: %v", err)
				continue
			}

			switch signal.Type {
			case "offer":
				log.Println("收到客户端的提议")
				sdp, err := base64.StdEncoding.DecodeString(signal.SDP)
				if err != nil {
					log.Printf("SDP 解码错误: %v", err)
					continue
				}
				log.Printf("解码的提议 SDP: %s", string(sdp)[:200]+"...")

				if err := peerConnection.SetRemoteDescription(webrtc.SessionDescription{
					Type: webrtc.SDPTypeOffer,
					SDP:  string(sdp),
				}); err != nil {
					log.Printf("设置远程描述错误: %v", err)
					continue
				}
				log.Println("成功设置远程描述")
				remoteDescSet = true

				for _, candidate := range pendingCandidates {
					if err := peerConnection.AddICECandidate(candidate); err != nil {
						log.Printf("添加缓存的 ICE 候选者错误: %v", err)
					} else {
						log.Println("成功添加缓存的 ICE 候选者")
					}
				}
				pendingCandidates = nil

				log.Println("创建应答")
				answer, err := peerConnection.CreateAnswer(nil)
				if err != nil {
					log.Printf("创建应答错误: %v", err)
					continue
				}

				log.Println("设置本地描述")
				if err := peerConnection.SetLocalDescription(answer); err != nil {
					log.Printf("设置本地描述错误: %v", err)
					continue
				}

				log.Printf("应答 SDP: %s", answer.SDP[:200]+"...")
				answerSDP := base64.StdEncoding.EncodeToString([]byte(answer.SDP))
				if err := conn.WriteJSON(Signal{
					SDP:  answerSDP,
					Type: "answer",
				}); err != nil {
					log.Printf("WebSocket 写入错误: %v", err)
					continue
				}
				log.Println("已发送应答给客户端")

				log.Println("等待 ICE 收集完成")
				gatherComplete := webrtc.GatheringCompletePromise(peerConnection)
				<-gatherComplete
				log.Println("ICE 收集完成")

				log.Println("启动 FFmpeg 视频流")
				go startFFmpegVideo(videoTrack)
				log.Println("启动 FFmpeg 音频流")
				go startFFmpegAudio(audioTrack)

			case "candidate":
				log.Println("收到客户端的 ICE 候选者")
				var candidate webrtc.ICECandidateInit
				if err := json.Unmarshal([]byte(signal.Candidate), &candidate); err != nil {
					log.Printf("ICE 候选者解析错误: %v", err)
					continue
				}
				if !remoteDescSet {
					log.Println("缓存 ICE 候选者，远程描述未设置")
					pendingCandidates = append(pendingCandidates, candidate)
				} else {
					if err := peerConnection.AddICECandidate(candidate); err != nil {
						log.Printf("添加 ICE 候选者错误: %v", err)
						continue
					}
					log.Println("成功添加 ICE 候选者")
				}
			default:
				log.Printf("未知消息类型: %s", signal.Type)
			}
		}
	}()

	wg.Wait()
}

func setupFFmpegLogger(cmd *exec.Cmd, trackType string) io.ReadCloser {
	cmd.Args = append(cmd.Args, "-loglevel", "info")
	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Printf("FFmpeg %s stderr 管道错误: %v", trackType, err)
		return nil
	}
	go func(stderr io.ReadCloser) {
		buf := make([]byte, 1024)
		for {
			n, err := stderr.Read(buf)
			if err != nil {
				if err != io.EOF {
					log.Printf("FFmpeg %s stderr 读取错误: %v", trackType, err)
				}
				return
			}
			log.Printf("FFmpeg %s stderr: %s", trackType, string(buf[:n]))
		}
	}(stderr)
	return stderr
}

func startFFmpegVideo(videoTrack *webrtc.TrackLocalStaticSample) {
	log.Println("准备 FFmpeg 视频命令")
	videoCmd := exec.Command("ffmpeg",
		"-re",
		"-i", "/home/weiqiangren/Documents/big-buck-bunny_trailer.mp4",
		"-c:v", "libx264",
		"-r", "25",
		"-crf", "23",
		"-b:v", "400k",
		"-maxrate", "450k",
		"-bufsize", "800k",
		"-g", "15",
		"-keyint_min", "15",
		"-sc_threshold", "0",
		"-preset", "fast",
		"-tune", "zerolatency",
		"-pix_fmt", "yuv420p",
		"-bsf:v", "h264_mp4toannexb",
		"-f", "h264",
		"pipe:1",
	)

	log.Println("创建 FFmpeg 视频 stderr 日志")
	setupFFmpegLogger(videoCmd, "video")

	log.Println("创建 FFmpeg 视频 stdout 管道")
	videoStdout, err := videoCmd.StdoutPipe()
	if err != nil {
		log.Fatalf("FFmpeg 视频 stdout 管道错误: %v", err)
	}

	log.Println("启动 FFmpeg 视频进程")
	if err := videoCmd.Start(); err != nil {
		log.Fatalf("FFmpeg 视频启动错误: %v", err)
	}

	go func() {
		log.Println("开始 FFmpeg 视频流")
		buf := make([]byte, 1024*1024)
		startTime := time.Now()
		sampleCount := uint64(0)

		for {
			n, err := videoStdout.Read(buf)
			if err != nil {
				if err == io.EOF {
					log.Println("FFmpeg 视频流结束 (EOF)")
					break
				}
				log.Printf("FFmpeg 视频读取错误: %v", err)
				break
			}

			if err := videoTrack.WriteSample(media.Sample{
				Data:     buf[:n],
				Duration: time.Second / 25,
			}); err != nil {
				log.Printf("视频轨道写入错误: %v", err)
				break
			}
			log.Printf("写入 %d 字节到 WebRTC 视频轨道", n)

			sampleCount++
			elapsed := time.Since(startTime)
			expected := time.Duration(sampleCount) * (time.Second / 25)
			if elapsed < expected {
				time.Sleep(expected - elapsed)
			}
		}
		videoStdout.Close()
	}()

	go func() {
		if err := videoCmd.Wait(); err != nil {
			log.Printf("FFmpeg 视频等待错误: %v", err)
		}
		log.Println("FFmpeg 视频进程退出")
	}()
}

func startFFmpegAudio(audioTrack *webrtc.TrackLocalStaticSample) {
	log.Println("准备 FFmpeg 音频命令")
	audioCmd := exec.Command("ffmpeg",
		"-re",
		"-i", "/home/weiqiangren/Documents/big-buck-bunny_trailer.mp4",
		"-map", "0:a:0",
		"-c:a", "libopus",
		"-b:a", "128k",
		"-vbr", "on",
		"-frame_duration", "20",
		"-ar", "48000",
		"-ac", "2",
		"-application", "voip", // 优化为 WebRTC 的 VoIP 场景
		"-f", "opus", // 直接输出 Opus 格式
		"pipe:1",
	)

	log.Println("创建 FFmpeg 音频 stderr 日志")
	setupFFmpegLogger(audioCmd, "audio")

	log.Println("创建 FFmpeg 音频 stdout 管道")
	audioStdout, err := audioCmd.StdoutPipe()
	if err != nil {
		log.Fatalf("FFmpeg 音频 stdout 管道错误: %v", err)
	}

	log.Println("启动 FFmpeg 音频进程")
	if err := audioCmd.Start(); err != nil {
		log.Fatalf("FFmpeg 音频启动错误: %v", err)
	}

	go func() {
		log.Println("开始 FFmpeg 音频流")
		buf := make([]byte, 1024*1024)
		startTime := time.Now()
		sampleCount := uint64(0)

		for {
			n, err := audioStdout.Read(buf)
			if err != nil {
				if err == io.EOF {
					log.Println("FFmpeg 音频流结束 (EOF)")
					break
				}
				log.Printf("FFmpeg 音频读取错误: %v", err)
				break
			}

			if err := audioTrack.WriteSample(media.Sample{
				Data:     buf[:n],
				Duration: time.Millisecond * 20,
			}); err != nil {
				log.Printf("音频轨道写入错误: %v", err)
				break
			}
			log.Printf("写入 %d 字节到 WebRTC 音频轨道", n)

			sampleCount++
			elapsed := time.Since(startTime)
			expected := time.Duration(sampleCount) * (time.Millisecond * 20)
			if elapsed < expected {
				time.Sleep(expected - elapsed)
			}
		}
		audioStdout.Close()
	}()

	go func() {
		if err := audioCmd.Wait(); err != nil {
			log.Printf("FFmpeg 音频等待错误: %v", err)
		}
		log.Println("FFmpeg 音频进程退出")
	}()
}
